import numpy as np
import struct
import matplotlib.pyplot as plt
from scipy.signal import stft
from scipy.fftpack import fftshift
import scipy

# 解决中文显示问题
plt.rcParams['font.sans-serif'] = ['KaiTi']  # 指定默认字体
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题

bin_file_path = r"C:\e\data\DATA5yi.bin"
fp = open(bin_file_path, "rb")  # 以二进制只读方式打开文件

output_bin = r"C:\e\data/abnormal_data(20m).bin"
fout = open(output_bin, "wb+")

N = int(20e6)
Fs = 100e6
shift = True  # 是否做fftshift
n = np.arange(0, N)
t = n / Fs

# F = 250e6
F = 200e6

x = 2 * np.pi * F * t
y = np.zeros([N], dtype=np.complex)
alpha = 10

y.real = alpha * np.cos(2 * np.pi * F * t)
y.imag = alpha * np.sin(2 * np.pi * F * t)

print("y.shape = ", y.shape)

iq_complex_cnt = int(N)
read_num = iq_complex_cnt * 2
read_bin_cnt = read_num * 2


i = np.arange(0, read_num, 2)
q = np.arange(1, read_num, 2)

content = fp.read(read_bin_cnt)
data = struct.unpack("h" * read_num, content)
data = np.array(data)

iq = np.zeros([iq_complex_cnt], dtype=np.complex)
iq.real = data[i]
iq.imag = data[q]

iq_add: np.ndarray = iq + y

print(iq_add.shape)

ab_iq = np.zeros_like(data)
ab_iq[i] = iq_add.real
ab_iq[q] = iq_add.imag
lst = ab_iq.tolist()
print(len(lst))
print(ab_iq.shape)
bin_data = struct.pack('h'*read_num, *lst)

fout.write(bin_data)


nfft = 2048
window = scipy.signal.get_window('hamming', Nx=2048)

F, T, Zxx = stft(y, fs=Fs, noverlap=1024, nfft=2048, window=window, return_onesided=False, nperseg=2048)
plt.figure(figsize=(10, 5))
s = np.abs(Zxx) / nfft
s = np.log10(s)
if shift:
    s = fftshift(s)
plt.imshow(s, aspect='auto', origin='lower')
plt.title("正弦信号")
plt.colorbar()
plt.show()


F, T, Zxx = stft(iq, fs=Fs, noverlap=1024, nfft=2048, window=window, return_onesided=False,
                 nperseg=2048)
plt.figure(figsize=(10, 5))
s = np.abs(Zxx) / nfft
s = np.log10(s)
if shift:
    s = fftshift(s)
plt.imshow(s, aspect='auto', origin='lower')
plt.title("正常信号")
plt.colorbar()
plt.show()


fout.flush()
fout.seek(0)
content = fout.read(read_bin_cnt)
data = struct.unpack("h" * read_num, content)
data = np.array(data)

iq = np.zeros([iq_complex_cnt], dtype=np.complex)
iq.real = data[i]
iq.imag = data[q]
F, T, Zxx = stft(iq, fs=Fs, noverlap=1024, nfft=2048, window=window, return_onesided=False,
                 nperseg=2048)
plt.figure(figsize=(10, 5))
s = np.abs(Zxx) / nfft
s = np.log10(s)
if shift:
    s = fftshift(s)
plt.imshow(s, aspect='auto', origin='lower')
plt.title("二进制转换完的异常信号")
plt.colorbar()
plt.show()


fout.close()
fp.close()
print()
